/**
 * @file drm_device.c
 * @brief DRM display handling implementation with Atomic KMS API.
 *
 * Implements functions for initializing a DRM device, finding resources,
 * and displaying DMABUF-backed framebuffers using non-blocking atomic commits
 * for tear-free, correctly scaled video presentation.
 */

#include <unistd.h>
#include <sys/ioctl.h>
#include <stdlib.h>
#include <string.h>
#include <poll.h>

#include "drm_device.h"

// --- Page Flip Event Handler ---
static void page_flip_handler(int fd, unsigned int frame, unsigned int sec, unsigned int usec, void *data)
{
    (void)fd; (void)frame; (void)sec; (void)usec;
    drm_device_t *ctx = (drm_device_t *)data;
    ctx->page_flip_pending = false;
}

// --- Internal Helper Functions ---

// Helper to find a specific DRM property ID for a given object.
static int find_property_id(int drmfd, uint32_t object_id, uint32_t object_type, const char *prop_name, uint32_t *prop_id)
{
    drmModeObjectPropertiesPtr props = drmModeObjectGetProperties(drmfd, object_id, object_type);
    if (!props) {
        fprintf(stderr, "Failed to get properties for object %u\n", object_id);
        return -1;
    }

    for (uint32_t i = 0; i < props->count_props; i++)
    {
        drmModePropertyPtr prop = drmModeGetProperty(drmfd, props->props[i]);
        if (prop)
        {
            if (strcmp(prop->name, prop_name) == 0)
            {
                *prop_id = prop->prop_id;
                drmModeFreeProperty(prop);
                drmModeFreeObjectProperties(props);
                return 0;
            }
            drmModeFreeProperty(prop);
        }
    }
    fprintf(stderr, "Could not find property '%s' for object %u\n", prop_name, object_id);
    drmModeFreeObjectProperties(props);
    return -1;
}


static int drm_find_resources_and_props(int drmfd, drm_device_t *ctx)
{
    drmModeRes *res = drmModeGetResources(drmfd);
    RETURN_ON_FAIL(!res, "drmModeGetResources failed: %s\n", ERRSTR);

    // Find a connected connector
    for (int i = 0; i < res->count_connectors; i++)
    {
        drmModeConnector *connector = drmModeGetConnector(drmfd, res->connectors[i]);
        if (connector && connector->connection == DRM_MODE_CONNECTED)
        {
            ctx->connector_id = connector->connector_id;
            
            char target_mode_str[32];
            snprintf(target_mode_str, sizeof(target_mode_str), "%dx%d", ctx->width, ctx->height);
            for (int j = 0; j < connector->count_modes; j++) {
                if (strcmp(connector->modes[j].name, target_mode_str) == 0) {
                    memcpy(&ctx->mode, &connector->modes[j], sizeof(drmModeModeInfo));
                    goto found_mode;
                }
            }
            // Fallback to preferred mode
            memcpy(&ctx->mode, &connector->modes[0], sizeof(drmModeModeInfo));
            printf("WARN: No exact mode match for %dx%d. Falling back to %s\n", ctx->width, ctx->height, ctx->mode.name);

        found_mode:
            drmModeEncoder *encoder = drmModeGetEncoder(drmfd, connector->encoder_id);
            if (encoder)
            {
                for(int k=0; k < res->count_crtcs; ++k) {
                    if (encoder->possible_crtcs & (1 << k)) {
                       ctx->crtc_id = res->crtcs[k];
                       ctx->crtc_idx = k;
                       break;
                    }
                }
                drmModeFreeEncoder(encoder);
            }
            drmModeFreeConnector(connector);
            break;
        }
        drmModeFreeConnector(connector);
    }
    RETURN_ON_FAIL(ctx->crtc_id == 0, "Could not find a valid CRTC for a connected connector\n");

    // Find a suitable plane
    drmModePlaneResPtr planes = drmModeGetPlaneResources(drmfd);
    RETURN_ON_FAIL(!planes, "drmModeGetPlaneResources failed: %s\n", ERRSTR);

    for (unsigned int i = 0; i < planes->count_planes; i++)
    {
        drmModePlanePtr plane = drmModeGetPlane(drmfd, planes->planes[i]);
        if (!plane) continue;

        if (plane->possible_crtcs & (1 << ctx->crtc_idx)) {
             for (unsigned int j = 0; j < plane->count_formats; j++) {
                if (plane->formats[j] == ctx->out_fourcc) {
                    ctx->plane_id = plane->plane_id;
                    break;
                }
            }
        }
        drmModeFreePlane(plane);
        if (ctx->plane_id) break;
    }
    drmModeFreePlaneResources(planes);
    RETURN_ON_FAIL(ctx->plane_id == 0, "Failed to find a suitable plane\n");

    // Find property IDs for atomic requests
    uint32_t prop_crtc_active, prop_crtc_mode_id, prop_conn_crtc_id;
    RETURN_ON_FAIL(find_property_id(drmfd, ctx->crtc_id, DRM_MODE_OBJECT_CRTC, "ACTIVE", &prop_crtc_active), "Failed to find ACTIVE prop for CRTC\n");
    RETURN_ON_FAIL(find_property_id(drmfd, ctx->crtc_id, DRM_MODE_OBJECT_CRTC, "MODE_ID", &prop_crtc_mode_id), "Failed to find MODE_ID prop for CRTC\n");
    RETURN_ON_FAIL(find_property_id(drmfd, ctx->connector_id, DRM_MODE_OBJECT_CONNECTOR, "CRTC_ID", &prop_conn_crtc_id), "Failed to find CRTC_ID prop for Connector\n");

    RETURN_ON_FAIL(find_property_id(drmfd, ctx->plane_id, DRM_MODE_OBJECT_PLANE, "FB_ID", &ctx->prop_fb_id), "Could not find FB_ID property\n");
    RETURN_ON_FAIL(find_property_id(drmfd, ctx->plane_id, DRM_MODE_OBJECT_PLANE, "CRTC_ID", &ctx->prop_crtc_id), "Could not find CRTC_ID property\n");
    RETURN_ON_FAIL(find_property_id(drmfd, ctx->plane_id, DRM_MODE_OBJECT_PLANE, "SRC_X", &ctx->prop_src_x), "Could not find SRC_X property\n");
    RETURN_ON_FAIL(find_property_id(drmfd, ctx->plane_id, DRM_MODE_OBJECT_PLANE, "SRC_Y", &ctx->prop_src_y), "Could not find SRC_Y property\n");
    RETURN_ON_FAIL(find_property_id(drmfd, ctx->plane_id, DRM_MODE_OBJECT_PLANE, "SRC_W", &ctx->prop_src_w), "Could not find SRC_W property\n");
    RETURN_ON_FAIL(find_property_id(drmfd, ctx->plane_id, DRM_MODE_OBJECT_PLANE, "SRC_H", &ctx->prop_src_h), "Could not find SRC_H property\n");
    RETURN_ON_FAIL(find_property_id(drmfd, ctx->plane_id, DRM_MODE_OBJECT_PLANE, "CRTC_X", &ctx->prop_crtc_x), "Could not find CRTC_X property\n");
    RETURN_ON_FAIL(find_property_id(drmfd, ctx->plane_id, DRM_MODE_OBJECT_PLANE, "CRTC_Y", &ctx->prop_crtc_y), "Could not find CRTC_Y property\n");
    RETURN_ON_FAIL(find_property_id(drmfd, ctx->plane_id, DRM_MODE_OBJECT_PLANE, "CRTC_W", &ctx->prop_crtc_w), "Could not find CRTC_W property\n");
    RETURN_ON_FAIL(find_property_id(drmfd, ctx->plane_id, DRM_MODE_OBJECT_PLANE, "CRTC_H", &ctx->prop_crtc_h), "Could not find CRTC_H property\n");
    
    printf("DRM resources found: CRTC=%u, Plane=%u, Connector=%u\n", ctx->crtc_id, ctx->plane_id, ctx->connector_id);
    printf("DRM Mode Set: %s (%dx%d@%d)\n", ctx->mode.name, ctx->mode.hdisplay, ctx->mode.vdisplay, ctx->mode.vrefresh);
    
    drmModeFreeResources(res);
    return 0;
}


// --- Public API Implementation ---

int drm_init(drm_device_t *ctx)
{
    ctx->fd = drmOpen(ctx->module, NULL);
    RETURN_ON_FAIL(ctx->fd < 0, "drmOpen(%s) failed: %s\n", ctx->module, ERRSTR);

    int ret = drmSetClientCap(ctx->fd, DRM_CLIENT_CAP_ATOMIC, 1);
    RETURN_ON_FAIL(ret, "Failed to set atomic cap: %s\n", ERRSTR);

    ret = drm_find_resources_and_props(ctx->fd, ctx);
    RETURN_ON_FAIL(ret, "Failed to find DRM resources\n");

    memset(&ctx->ev_ctx, 0, sizeof(drmEventContext));
    ctx->ev_ctx.version = DRM_EVENT_CONTEXT_VERSION;
    ctx->ev_ctx.page_flip_handler = page_flip_handler;
    ctx->page_flip_pending = false;

    return 0;
}

void drm_cleanup(drm_device_t *ctx)
{
    if (ctx->fd < 0) return;

    if (ctx->buffers)
    {
        for (unsigned int i = 0; i < ctx->buffer_count; i++)
        {
            if (ctx->buffers[i].fb_handle)
                drmModeRmFB(ctx->fd, ctx->buffers[i].fb_handle);

            if (ctx->buffers[i].bo_handle)
            {
                struct drm_mode_destroy_dumb gem_destroy = {.handle = ctx->buffers[i].bo_handle};
                ioctl(ctx->fd, DRM_IOCTL_MODE_DESTROY_DUMB, &gem_destroy);
            }
            if (ctx->buffers[i].dma_fd >= 0)
                close(ctx->buffers[i].dma_fd);
        }
        free(ctx->buffers);
        ctx->buffers = NULL;
    }

    drmClose(ctx->fd);
    ctx->fd = -1;
}

int drm_create_buffers(drm_device_t *ctx, int width, int height)
{
    ctx->buffer_width = width;
    ctx->buffer_height = height;

    ctx->buffers = calloc(ctx->buffer_count, sizeof(struct drm_buffer));
    RETURN_ON_FAIL(!ctx->buffers, "Failed to allocate buffer array\n");

    for (unsigned int i = 0; i < ctx->buffer_count; ++i)
    {
        struct drm_buffer *b = &ctx->buffers[i];
        b->dma_fd = -1;

        struct drm_mode_create_dumb gem = {
            .width = (uint32_t)ctx->buffer_width,
            .height = (uint32_t)ctx->buffer_height,
            .bpp = 12, 
        };
        int ret = ioctl(ctx->fd, DRM_IOCTL_MODE_CREATE_DUMB, &gem);
        RETURN_ON_FAIL(ret, "DRM_IOCTL_MODE_CREATE_DUMB failed: %s\n", ERRSTR);
        b->bo_handle = gem.handle;

        uint32_t pitches[4] = {(uint32_t)ctx->buffer_width, (uint32_t)ctx->buffer_width};
        uint32_t offsets[4] = {0, pitches[0] * (uint32_t)ctx->buffer_height};
        uint32_t bo_handles[4] = {b->bo_handle, b->bo_handle};

        ret = drmModeAddFB2(ctx->fd, (uint32_t)ctx->buffer_width, (uint32_t)ctx->buffer_height, ctx->out_fourcc,
                            bo_handles, pitches, offsets, &b->fb_handle, 0);
        RETURN_ON_FAIL(ret, "drmModeAddFB2 failed: %s\n", ERRSTR);

        struct drm_prime_handle prime = { .handle = b->bo_handle };
        ret = ioctl(ctx->fd, DRM_IOCTL_PRIME_HANDLE_TO_FD, &prime);
        RETURN_ON_FAIL(ret, "DRM_IOCTL_PRIME_HANDLE_TO_FD failed: %s\n", ERRSTR);
        b->dma_fd = prime.fd;
    }
    return 0;
}

int drm_set_plane(drm_device_t *ctx, int buffer_index)
{
    if (buffer_index < 0 || (unsigned int)buffer_index >= ctx->buffer_count) {
        return -1;
    }
    struct drm_buffer *b = &ctx->buffers[buffer_index];
    drmModeAtomicReq *req = drmModeAtomicAlloc();
    RETURN_ON_FAIL(!req, "Failed to allocate atomic request\n");

    uint32_t blob_id;
    drmModeCreatePropertyBlob(ctx->fd, &ctx->mode, sizeof(ctx->mode), &blob_id);

    // --- This is a full modeset, linking connector, crtc, and plane. ---
    uint32_t prop_active, prop_mode_id, prop_conn_crtc_id;
    find_property_id(ctx->fd, ctx->crtc_id, DRM_MODE_OBJECT_CRTC, "ACTIVE", &prop_active);
    find_property_id(ctx->fd, ctx->crtc_id, DRM_MODE_OBJECT_CRTC, "MODE_ID", &prop_mode_id);
    find_property_id(ctx->fd, ctx->connector_id, DRM_MODE_OBJECT_CONNECTOR, "CRTC_ID", &prop_conn_crtc_id);
    
    drmModeAtomicAddProperty(req, ctx->connector_id, prop_conn_crtc_id, ctx->crtc_id);
    drmModeAtomicAddProperty(req, ctx->crtc_id, prop_active, 1);
    drmModeAtomicAddProperty(req, ctx->crtc_id, prop_mode_id, blob_id);
    
    // --- Configure plane for scaling ---
    drmModeAtomicAddProperty(req, ctx->plane_id, ctx->prop_crtc_id, ctx->crtc_id);
    drmModeAtomicAddProperty(req, ctx->plane_id, ctx->prop_fb_id, b->fb_handle);
    drmModeAtomicAddProperty(req, ctx->plane_id, ctx->prop_src_x, 0); // Source rect: full buffer
    drmModeAtomicAddProperty(req, ctx->plane_id, ctx->prop_src_y, 0);
    drmModeAtomicAddProperty(req, ctx->plane_id, ctx->prop_src_w, (uint64_t)ctx->buffer_width << 16);
    drmModeAtomicAddProperty(req, ctx->plane_id, ctx->prop_src_h, (uint64_t)ctx->buffer_height << 16);
    drmModeAtomicAddProperty(req, ctx->plane_id, ctx->prop_crtc_x, 0); // Destination rect: full screen
    drmModeAtomicAddProperty(req, ctx->plane_id, ctx->prop_crtc_y, 0);
    drmModeAtomicAddProperty(req, ctx->plane_id, ctx->prop_crtc_w, ctx->mode.hdisplay);
    drmModeAtomicAddProperty(req, ctx->plane_id, ctx->prop_crtc_h, ctx->mode.vdisplay);

    // Blocking commit for the first frame
    int ret = drmModeAtomicCommit(ctx->fd, req, DRM_MODE_ATOMIC_ALLOW_MODESET, NULL);
    drmModeAtomicFree(req);
    drmModeDestroyPropertyBlob(ctx->fd, blob_id);
    RETURN_ON_FAIL(ret, "Initial atomic commit failed: %s\n", ERRSTR);

    return 0;
}

int drm_page_flip(drm_device_t *ctx, int buffer_index)
{
    if (buffer_index < 0 || (unsigned int)buffer_index >= ctx->buffer_count)
    {
        return -1;
    }
    struct drm_buffer *b = &ctx->buffers[buffer_index];
    drmModeAtomicReq *req = drmModeAtomicAlloc();
    RETURN_ON_FAIL(!req, "Failed to allocate atomic request\n");

    // For a simple flip, we only need to update the framebuffer ID on the plane.
    // The scaling properties (SRC/CRTC rects) were set in drm_set_plane and persist.
    // However, explicitly setting them again is safer.
    int ret_prop = 0;
    ret_prop |= drmModeAtomicAddProperty(req, ctx->plane_id, ctx->prop_fb_id, b->fb_handle);
    ret_prop |= drmModeAtomicAddProperty(req, ctx->plane_id, ctx->prop_crtc_id, ctx->crtc_id);
    ret_prop |= drmModeAtomicAddProperty(req, ctx->plane_id, ctx->prop_src_x, 0);
    ret_prop |= drmModeAtomicAddProperty(req, ctx->plane_id, ctx->prop_src_y, 0);
    ret_prop |= drmModeAtomicAddProperty(req, ctx->plane_id, ctx->prop_src_w, (uint64_t)ctx->buffer_width << 16);
    ret_prop |= drmModeAtomicAddProperty(req, ctx->plane_id, ctx->prop_src_h, (uint64_t)ctx->buffer_height << 16);
    ret_prop |= drmModeAtomicAddProperty(req, ctx->plane_id, ctx->prop_crtc_x, 0);
    ret_prop |= drmModeAtomicAddProperty(req, ctx->plane_id, ctx->prop_crtc_y, 0);
    ret_prop |= drmModeAtomicAddProperty(req, ctx->plane_id, ctx->prop_crtc_w, ctx->mode.hdisplay);
    ret_prop |= drmModeAtomicAddProperty(req, ctx->plane_id, ctx->prop_crtc_h, ctx->mode.vdisplay);

    if (ret_prop < 0) {
        fprintf(stderr, "Failed to add property to atomic request\n");
        drmModeAtomicFree(req);
        return -1;
    }
    
    // Non-blocking commit with page flip event
    int ret = drmModeAtomicCommit(ctx->fd, req, DRM_MODE_PAGE_FLIP_EVENT | DRM_MODE_ATOMIC_NONBLOCK, ctx);
    drmModeAtomicFree(req);
    RETURN_ON_FAIL(ret, "Atomic page flip failed: %s\n", ERRSTR);
    
    ctx->page_flip_pending = true;
    return 0;
}

int drm_wait_for_flip(drm_device_t *ctx)
{
    struct pollfd p = {.fd = ctx->fd, .events = POLLIN};
    
    while (ctx->page_flip_pending)
    {
        int ret = poll(&p, 1, 2000); // 2-second timeout
        if (ret < 0)
        {
            if (errno == EINTR) continue;
            fprintf(stderr, "DRM poll failed: %s\n", ERRSTR);
            return -1;
        }
        if (ret == 0)
        {
            fprintf(stderr, "DRM poll timeout.\n");
            ctx->page_flip_pending = false; 
            return -1;
        }
        drmHandleEvent(ctx->fd, &ctx->ev_ctx);
    }
    return 0;
}

